# wider2yolo_fixpath.py
import os, cv2, shutil
from pathlib import Path

# ========== 1. 路径写死 ==========
ROOT      = Path('data/WIDER')          # 数据根目录
IMG_TRAIN = ROOT / 'WIDER_train/images'               # 原始训练图片
IMG_VAL   = ROOT / 'WIDER_val/images'                 # 原始验证图片
GT_TRAIN  = ROOT / 'wider_face_split/wider_face_train_bbx_gt.txt'
GT_VAL    = ROOT / 'wider_face_split/wider_face_val_bbx_gt.txt'

OUT_TRAIN = ROOT / 'train_yolo'                 # 输出目录
OUT_VAL   = ROOT / 'val_yolo'

# ========== 2. 转换函数 ==========
def convert(gt_txt, img_dir, out_dir):
    out_img = out_dir / 'images'
    out_lbl = out_dir / 'labels'
    out_img.mkdir(parents=True, exist_ok=True)
    out_lbl.mkdir(parents=True, exist_ok=True)

    with open(gt_txt) as f:
        lines = [x.rstrip() for x in f]

    idx = 0
    while idx < len(lines):
        img_name = lines[idx]; idx += 1          # 如 "0--Parade/0_Parade_xxx.jpg"
        n = int(lines[idx]); idx += 1

        boxes = []
        for _ in range(n):
            parts = lines[idx].split(); idx += 1
            x1, y1, w, h = map(float, parts[:4])
            if w <= 0 or h <= 0:
                continue
            boxes.append([x1, y1, w, h])

        # 复制图片
        src = img_dir / img_name
        dst = out_img / src.name
        dst.parent.mkdir(parents=True, exist_ok=True)
        shutil.copy(src, dst)

        # 写 YOLO 标签
        img = cv2.imread(str(src))
        h0, w0 = img.shape[:2]
        with open(out_lbl / f"{src.stem}.txt", 'w') as f:
            for x1, y1, w, h in boxes:
                xc = (x1 + w / 2) / w0
                yc = (y1 + h / 2) / h0
                wn = w / w0
                hn = h / h0
                f.write(f"0 {xc:.6f} {yc:.6f} {wn:.6f} {hn:.6f}\n")

# ========== 3. 一键执行 ==========
if __name__ == '__main__':
    convert(GT_TRAIN, IMG_TRAIN, OUT_TRAIN)
    convert(GT_VAL,   IMG_VAL,   OUT_VAL)
    print('✅  train & val 均已转好！')